# 导入logging模块，用于日志记录
import logging

# 导入json模块，用于处理JSON数据
import json

# 导入uuid模块，用于生成唯一会话ID
import uuid

# 导入time模块，用于时间戳
import time

# 从tools模块导入tools对象
from tools import tools, completion_data

# 导入datetime模块，用于处理日期和时间
import datetime

# 从http.server模块导入ThreadingHTTPServer和BaseHTTPRequestHandler类
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler

# 配置日志记录级别为INFO
logging.basicConfig(level=logging.INFO)

# 获取当前模块的日志记录器
logger = logging.getLogger(__name__)


# 定义MCPServer类，作为MCP服务器的核心类
class MCPServer:
    """MCP服务器核心类"""

    # 初始化方法
    def __init__(self):
        # 协议版本号
        self.protocol_version = "2025-06-18"
        # 服务器信息，包括名称和版本
        self.server_info = {"name": "MCP服务器", "version": "1.0.0"}
        # 服务器能力，包括工具、资源和提示
        self.capabilities = {
            "tools": {},
            "resources": {},
            "prompts": {},
            "completions": {},
        }
        # 会话管理字典
        self.sessions = {}

    def _get_language_completions(self, partial_value):
        if not partial_value:
            return completion_data["languages"]
        return [
            lang
            for lang in completion_data["languages"]
            if partial_value.lower() in lang.lower()
        ]

    def _get_framework_completions(self, language, partial_value):
        frameworks = completion_data["frameworks"].get(language, [])
        if not partial_value:
            return frameworks
        return [
            framework
            for framework in frameworks
            if partial_value.lower() in framework.lower()
        ]

    def _get_all_framework_completions(self, partial_value):
        all_frameworks = []
        for frameworks in completion_data["frameworks"].values():
            all_frameworks.extend(frameworks)
        if not partial_value:
            return all_frameworks
        return [
            framework
            for framework in all_frameworks
            if partial_value.lower() in framework.lower()
        ]

    def _get_path_completions(self, partial_value):
        files = ["README.md", "main.py", "UPDATE.md"]
        if not partial_value:
            return files
        return [file for file in files if partial_value.lower() in file.lower()]

    def _get_analysis_type_completions(self, partial_value):
        if not partial_value:
            return completion_data["analysis_types"]
        return [
            analysis_type
            for analysis_type in completion_data["analysis_types"]
            if partial_value.lower() in analysis_type.lower()
        ]

    def handle_completion_complete(self, params, session_id):
        if session_id:
            self._get_or_create_session(session_id)
        ref = params.get("ref")
        context = params.get("context", {})
        ref_type = ref.get("type")

        argument = params.get("argument")
        argument_name = argument.get("name")
        argument_value = argument.get("value")
        # 根据引用类型和参数名称提供相应的补全的建议
        completion_values = []
        if ref_type == "ref/prompt":
            ref_name = ref.get("name")
            if ref_name == "code_review":
                if argument_name == "language":
                    completion_values = self._get_language_completions(argument_value)
                elif argument_name == "framework":
                    selected_language = context.get("arguments", {}).get("language")
                    if selected_language:
                        completion_values = self._get_framework_completions(
                            selected_language, argument_value
                        )
                    else:
                        completion_values = self._get_all_framework_completions(
                            argument_value
                        )
        elif ref_type == "ref/resource":
            uri = ref.get("uri")
            print(uri)
            print(argument_value)
            print(argument_name)
            if "file:///" in uri:
                if argument_name == "analysis_type":
                    completion_values = self._get_analysis_type_completions(
                        argument_value
                    )
                elif argument_name == "path":
                    completion_values = self._get_path_completions(argument_value)
        if len(completion_values) > 100:
            completion_values = completion_values[:100]
        return {
            "completion": {
                "values": completion_values,
                "total": len(completion_values),
                "hasMore": False,
            }
        }

    # 获取所有活跃会话信息
    def get_all_sessions(self):
        """获取所有活跃会话信息"""
        return self.sessions.copy()

    # 处理工具列表请求
    def handle_tools_list(self, params=None, session_id: str = None):
        """处理工具列表请求"""
        # 如果有会话ID，则获取或创建会话
        if session_id:
            self._get_or_create_session(session_id)
        # 返回工具列表
        return {"tools": tools}

    # 处理初始化请求（旧方法，已被覆盖）
    def handle_initialize(self, params, session_id: str = None):
        """处理初始化请求"""
        # 返回服务器能力
        return {"capabilities": self.capabilities}

    # 生成唯一的会话ID
    def _generate_session_id(self) -> str:
        """生成唯一的会话ID"""
        # 使用uuid4生成唯一ID
        return str(uuid.uuid4())

    # 获取会话信息
    def get_session_info(self, session_id: str):
        """获取会话信息"""
        # 如果会话ID存在于sessions字典中，返回会话信息
        if session_id in self.sessions:
            return self.sessions[session_id]
        # 否则返回None
        return None

    def _cleanup_expired_sessions(self, max_age: int = 3600):
        """清理过期会话（默认1小时）"""
        current_time = time.time()
        expired_sessions = []

        for session_id, session_data in self.sessions.items():
            if current_time - session_data["last_activity"] > max_age:
                expired_sessions.append(session_id)

        for session_id in expired_sessions:
            del self.sessions[session_id]
            logger.info(f"清理过期会话: {session_id}")

        return len(expired_sessions)

    # 获取或创建会话
    def _get_or_create_session(self, session_id: str = None):
        """获取或创建会话"""
        # 如果没有传入会话ID，则生成一个新的
        if not session_id:
            session_id = self._generate_session_id()
        # 如果会话ID不存在于sessions字典中，则新建一个会话
        if session_id not in self.sessions:
            self.sessions[session_id] = {
                "created_at": time.time(),
                "last_activity": time.time(),
                "message_count": 0,
                "client_info": None,
            }
            # 记录新会话的创建
            logger.info(f"创建新会话: {session_id}")
        # 更新会话的最后活动时间
        self.sessions[session_id]["last_activity"] = time.time()
        # 增加消息计数
        self.sessions[session_id]["message_count"] += 1
        # 返回会话ID
        return session_id

    # 处理工具调用请求
    def handle_tools_call(self, params, session_id: str = None):
        """处理工具调用请求"""
        # 如果有会话ID，则获取或创建会话
        if session_id:
            self._get_or_create_session(session_id)

        # 获取调用列表
        calls = params.get("calls", [])
        # 初始化结果列表
        results = []

        # 遍历每个调用
        for call in calls:
            # 获取工具名称
            tool_name = call.get("name")
            # 获取参数
            arguments = call.get("arguments", {})

            # 如果工具名称为"calculate"
            if tool_name == "calculate":
                # 获取表达式
                expression = arguments.get("expression", "")
                try:
                    # 计算表达式结果
                    result = eval(expression)
                    # 添加成功结果到结果列表
                    results.append(
                        {
                            "name": tool_name,
                            "content": [
                                {
                                    "type": "text",
                                    "text": f"计算结果: {expression} = {result}",
                                }
                            ],
                        }
                    )
                except Exception as e:
                    # 计算出错，添加错误信息
                    results.append(
                        {
                            "name": tool_name,
                            "isError": True,
                            "error": {"message": f"计算错误: {str(e)}"},
                        }
                    )
            elif tool_name == "code_review":
                results.append({"name": tool_name})
            elif tool_name == "file_analyzer":
                results.append(
                    {
                        "name": tool_name,
                    }
                )
            else:
                # 未知工具，添加错误信息
                results.append(
                    {
                        "name": tool_name,
                        "isError": True,
                        "error": {"message": f"未知工具: {tool_name}"},
                    }
                )

        # 返回所有调用的结果
        return {"calls": results}

    # 处理初始化请求（实际生效的方法，覆盖上面的方法）
    def handle_initialize(self, params, session_id: str = None):
        """处理初始化请求"""
        # 获取客户端信息
        client_info = params.get("clientInfo", {})
        # 获取或创建会话ID
        session_id = self._get_or_create_session(session_id)
        # 存储客户端信息到会话
        self.sessions[session_id]["client_info"] = client_info
        # 记录客户端初始化信息
        logger.info(f"客户端初始化: {client_info}, 会话ID: {session_id}")
        # 返回初始化结果，包括协议版本、能力、服务器信息和会话ID
        return {
            "protocolVersion": self.protocol_version,
            "capabilities": self.capabilities,
            "serverInfo": self.server_info,
            "sessionId": session_id,  # 返回会话ID给客户端
        }


# 定义MCPHTTPHandler类，继承自BaseHTTPRequestHandler
class MCPHTTPHandler(BaseHTTPRequestHandler):
    """HTTP处理器，支持Streamable HTTP传输"""

    # 初始化方法，接收mcp_server参数
    def __init__(self, *args, mcp_server=None, **kwargs):
        # 保存MCP服务器实例
        self.mcp_server = mcp_server
        # 调用父类的初始化方法
        super().__init__(*args, **kwargs)

    # 处理GET请求（建立SSE流）
    def do_GET(self):
        """处理GET请求（建立SSE流）"""
        try:
            # 如果请求路径不是/mcp，返回404
            if self.path != "/mcp":
                self.send_error(404, "MCP endpoint not found")
                return

            # 检查Accept头
            accept_header = self.headers.get("Accept", "")
            # 如果Accept头不包含text/event-stream，返回405
            if "text/event-stream" not in accept_header:
                self.send_error(405, "Method not allowed")
                return

            # 提取会话ID
            session_id = self.headers.get("Mcp-Session-Id")
            if session_id:
                # 验证会话是否存在
                session_info = self.mcp_server.get_session_info(session_id)
                if not session_info:
                    self.send_error(401, "Invalid session ID")
                    return
                # 记录使用现有会话
                logger.info(f"SSE连接使用现有会话: {session_id}")
            else:
                # 为SSE连接创建新会话
                session_id = self.mcp_server._get_or_create_session(None)
                # 记录新会话
                logger.info(f"SSE连接创建新会话: {session_id}")

            # 设置SSE头
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.send_header("Connection", "keep-alive")
            self.send_header("MCP-Protocol-Version", self.mcp_server.protocol_version)
            self.send_header("Mcp-Session-Id", session_id)  # 返回会话ID
            self.end_headers()

            # 发送初始事件
            self._send_sse_event(
                "connected", {"status": "connected", "sessionId": session_id}
            )
            # 记录SSE连接建立
            logger.info(f"SSE连接已建立，会话ID: {session_id}，开始发送时间更新")

            # 定时发送服务器时间
            try:
                while True:
                    # 每5秒发送一次
                    time.sleep(5)
                    # 获取当前服务器时间
                    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    # 发送时间更新事件
                    self._send_sse_event(
                        "time_update",
                        {
                            "server_time": current_time,
                            "timestamp": time.time(),
                            "sessionId": session_id,
                        },
                    )
                    # 记录发送时间
                    logger.info(f"发送时间更新: {current_time}, 会话ID: {session_id}")
            except (BrokenPipeError, ConnectionResetError):
                # 连接断开，记录日志
                logger.info(f"SSE连接已断开，会话ID: {session_id}")
            except Exception as e:
                # 发送事件出错，记录日志
                logger.error(f"SSE事件发送错误: {e}, 会话ID: {session_id}")

        except Exception as e:
            # 处理GET请求出错，记录日志
            logger.error(f"GET处理错误: {e}")
            # 返回500错误
            self.send_error(500, "Internal Server Error")

    # 发送SSE事件
    def _send_sse_event(self, event_type: str, data):
        """发送SSE事件"""
        try:
            # 构造事件数据
            event_data = f"event: {event_type}\n"
            event_data += f"data: {json.dumps(data)}\n\n"
            # 写入事件到输出流
            self.wfile.write(event_data.encode("utf-8"))
            self.wfile.flush()
        except Exception as e:
            # 发送事件出错，记录日志
            logger.error(f"发送SSE事件错误: {e}")

    # 处理MCP消息
    def _handle_mcp_message(self, message, session_id: str = None):
        """处理MCP消息"""
        # 获取方法名
        method = message.get("method")
        # 获取参数
        params = message.get("params", {})
        # 获取消息ID
        msg_id = message.get("id")
        # 判断方法类型
        if method == "initialize":
            # 调用服务器的handle_initialize方法
            result = self.mcp_server.handle_initialize(params, session_id)
            # 如果有消息ID，返回带ID的结果
            return {"id": msg_id, "result": result} if msg_id else None
        elif method == "tools/list":
            # 处理工具列表请求
            result = self.mcp_server.handle_tools_list(params, session_id)
            return {"id": msg_id, "result": result} if msg_id else None
        elif method == "tools/call":
            # 处理工具调用请求
            result = self.mcp_server.handle_tools_call(params, session_id)
            return {"id": msg_id, "result": result} if msg_id else None
        elif method == "completion/complete":
            # 处理工具调用请求
            result = self.mcp_server.handle_completion_complete(params, session_id)
            return {"id": msg_id, "result": result} if msg_id else None
        else:
            # 如果方法不存在，返回错误
            if msg_id:
                return {
                    "id": msg_id,
                    "error": {"code": -32601, "message": f"Method not found: {method}"},
                }
            # 没有ID则返回None
            return None

    # 处理POST请求（发送消息到服务器）
    def do_POST(self):
        """处理POST请求（发送消息到服务器）"""
        try:
            # 检查请求路径是否为/mcp
            if self.path != "/mcp":
                self.send_error(404, "MCP endpoint not found")
                return
            # 检查Accept头是否包含application/json或text/event-stream
            accept_header = self.headers.get("Accept", "")
            if (
                "application/json" not in accept_header
                and "text/event-stream" not in accept_header
            ):
                self.send_error(400, "Missing required Accept header")
                return
            # 检查MCP-Protocol-Version头是否匹配
            protocol_version = self.headers.get("MCP-Protocol-Version")
            if (
                protocol_version
                and protocol_version != self.mcp_server.protocol_version
            ):
                self.send_error(
                    400, f"Unsupported protocol version: {protocol_version}"
                )
                return
            # 获取会话ID
            session_id = self.headers.get("Mcp-Session-Id")
            # 读取请求体长度
            content_length = int(self.headers.get("Content-Length", 0))
            # 读取请求体内容
            body = self.rfile.read(content_length)
            try:
                # 尝试解析JSON消息
                message = json.loads(body.decode("utf-8"))
            except json.JSONDecodeError:
                # JSON解析失败
                self.send_error(400, "Invalid JSON")
                return
            # 处理MCP消息
            response = self._handle_mcp_message(message, session_id)
            # 发送响应
            if response:
                # 发送200响应码
                self.send_response(200)
                # 设置响应头Content-Type为application/json
                self.send_header("Content-Type", "application/json")
                # 设置协议版本头
                self.send_header(
                    "MCP-Protocol-Version", self.mcp_server.protocol_version
                )
                # 如果响应中包含会话ID，则添加到响应头
                if "result" in response and "sessionId" in response["result"]:
                    self.send_header("Mcp-Session-Id", response["result"]["sessionId"])
                # 结束响应头
                self.end_headers()
                # 写入JSON响应体
                self.wfile.write(json.dumps(response).encode("utf-8"))
            else:
                # 没有响应内容，返回202 Accepted
                self.send_response(202)  # Accepted
                self.end_headers()
        except Exception as e:
            # 记录错误日志
            logger.error(f"POST处理错误: {e}")
            # 返回500内部服务器错误
            self.send_error(500, "Internal Server Error")


# 定义运行HTTP服务器的函数
def run_http_server(mcp_server: MCPServer, host: str = "localhost", port: int = 8000):
    """运行HTTP服务器"""

    # 定义处理器工厂函数，用于传递mcp_server实例
    def handler_factory(*args, **kwargs):
        return MCPHTTPHandler(*args, mcp_server=mcp_server, **kwargs)

    # 创建ThreadingHTTPServer实例，绑定主机和端口
    server = ThreadingHTTPServer((host, port), handler_factory)
    # 记录服务器启动信息
    logger.info(f"HTTP服务器运行在 http://{host}:{port}/mcp")
    # 记录会话管理功能启用信息
    logger.info("会话管理功能已启用，支持会话ID跟踪和自动清理")

    # 启动会话清理线程
    import threading

    # 定义会话清理函数
    def cleanup_sessions():
        while True:
            try:
                # 每5分钟清理一次
                time.sleep(300)
                # 调用会话清理方法
                cleaned_count = mcp_server._cleanup_expired_sessions()
                # 如果有清理的会话，记录日志
                if cleaned_count > 0:
                    logger.info(f"清理了 {cleaned_count} 个过期会话")
                # 记录当前活跃会话数
                active_sessions = len(mcp_server.get_all_sessions())
                logger.debug(f"当前活跃会话数: {active_sessions}")
            except Exception as e:
                # 清理出错，记录日志
                logger.error(f"会话清理错误: {e}")

    # 启动清理线程
    cleanup_thread = threading.Thread(target=cleanup_sessions, daemon=True)
    cleanup_thread.start()
    logger.info("会话清理线程已启动")
    try:
        # 启动服务器，进入循环监听请求
        server.serve_forever()
    except KeyboardInterrupt:
        # 捕获Ctrl+C，记录服务器停止
        logger.info("HTTP服务器已停止")
        logger.info("正在清理所有会话...")
        # 清理所有会话
        all_sessions = mcp_server.get_all_sessions()
        for session_id in list(all_sessions.keys()):
            del mcp_server.sessions[session_id]
        logger.info(f"已清理 {len(all_sessions)} 个会话")
    finally:
        # 关闭服务器
        server.server_close()


# 定义主函数
def main():
    """主函数"""
    # 导入argparse模块，用于解析命令行参数
    import argparse

    # 创建ArgumentParser对象，设置描述信息
    parser = argparse.ArgumentParser(description="MCP HTTP服务器")
    # 添加--host参数，指定服务器主机，默认localhost
    parser.add_argument(
        "--host", default="localhost", help="HTTP服务器主机 (默认: localhost)"
    )
    # 添加--port参数，指定服务器端口，默认8000
    parser.add_argument(
        "--port", type=int, default=8000, help="HTTP服务器端口 (默认: 8000)"
    )

    # 解析命令行参数
    args = parser.parse_args()

    # 创建MCP服务器实例
    mcp_server = MCPServer()

    # 运行HTTP服务器
    run_http_server(mcp_server, args.host, args.port)


# 判断是否为主程序入口
if __name__ == "__main__":
    # 调用主函数
    main()
